package com.xiaozhen.tcp.zookeeper;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.*;
import org.apache.zookeeper.ZooDefs.Ids;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;

/**
 * @Author: zhenweisheng
 * @Date: 2020/7/31 13:53
 */
public class ServiceRegistry {

    private static final Logger logger = LoggerFactory.getLogger(ServiceRegistry.class);

    private static final String ROOT_NODE = "/route";

    private CountDownLatch latch = new CountDownLatch(1);

    private String registryAddress;

    public ServiceRegistry(String registryAddress) {
        this.registryAddress = registryAddress;
    }

    public void register(String childNode) {
        if (childNode != null) {
            ZooKeeper zk = connectServer();
            if (zk != null) {
                try {
                    if (null == zk.exists(ROOT_NODE, false)) {
                        createRootNode(zk);
                    }


                } catch (KeeperException e) {
                    e.printStackTrace();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                createNode(zk, childNode);
            }
        }
    }

    /**
     * 连接 zookeeper 服务器
     *
     * @return
     */
    private ZooKeeper connectServer() {
        ZooKeeper zk = null;
        try {
            zk = new ZooKeeper(registryAddress, 2000, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    if (event.getState() == Event.KeeperState.SyncConnected) {
                        latch.countDown();
                    }
                }
            });
            latch.await();
        } catch (Exception e) {
            logger.error("", e);
        }
        return zk;
    }

    /**
     * 创建顶级节点
     *
     * @param zk
     */
    private void createRootNode(ZooKeeper zk) {
        try {
            if (null == zk.exists(ROOT_NODE, false)) {
                String path = zk.create(ROOT_NODE, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                logger.info("创建路由根节点 ({})", path);
            }
        } catch (KeeperException | InterruptedException e) {
            logger.error("创建路由根节点异常:", e);
        }
    }

    /**
     * 创建子节点
     *
     * @param zk
     * @param data
     */
    private void createNode(ZooKeeper zk, String data) {
        try {
            String node = ROOT_NODE + "/" + data;
            if (null == zk.exists(node, false)) {
                byte[] bytes = data.getBytes();
                String path = zk.create(node, bytes, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                logger.info("创建节点 ({} => {})", path, data);
            }
        } catch (KeeperException | InterruptedException e) {
            logger.error("创建节点异常:", e);
        }
    }


    /**
     * 如果使用单机版请注掉下面的代码，则模拟前置不会继续连接zookeeper
     */
    public void registerToZookeeper(String data) {
        RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 10);
        CuratorFramework cf = CuratorFrameworkFactory.builder()
                .connectString(registryAddress)
                .sessionTimeoutMs(6000)
                .retryPolicy(retryPolicy)
                .build();
        System.out.println("zk连接中。。。。。。");
        //3 开启连接
        cf.start();
        while (cf.getState() != CuratorFrameworkState.STARTED) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        System.out.println("zk连接成功。。。。。");


        try {
            cf.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL).forPath("/route/" + data, data.getBytes());
            System.out.println("********zookeeper注册前置信息成功！********");
        } catch (Exception e) {
            System.err.println("zookeeper注册前置信息失败");
        }
    }
}

